import os
import requests
import json
import pdfplumber
from docx import Document
import csv

# -------------------------- 配置参数 --------------------------
# TODO: 请替换为你的API Key
API_KEY = "41834af386ab4067a71e787fd958a9c6.kWaJQtYORQGXutw9"
# 智谱chat接口
API_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
# 待处理的文档文件夹路径
# TODO: 将此处替换为你需要处理的文档文件夹路径
DOCUMENTS_DIR = r"D:\云烟-md\总糖-md"
# 输出CSV文件的路径
OUTPUT_CSV_PATH = r"D:\Users\wj\PycharmProjects\信息提取\md-结构化数据\总糖.csv"
# 定义CSV文件的列头
COLUMN_HEADERS = [
    '甜味剂类型', '论文名称', '致甜成分名称', '致甜成分的优点',
    '致甜成分的缺点', '致甜成分实验方法', '评吸方法', '卷烟余味',
    '致甜成分的应用', '致甜成分的添加方法', '致甜成分添加量及口感',
    '甜味持久性', '致甜成分相对甜度', '致甜成分烟气烟丝中的含量及迁移率',
    '致甜成分分析方法', '致甜成分的香气', '致甜成分的添加部位',
    '烟叶类型'
]


# -------------------------- 工具函数 --------------------------
def read_document(file_path):
    """读取文档内容（支持PDF、DOCX、MD）"""
    content = ""
    try:
        if file_path.lower().endswith(".pdf"):
            with pdfplumber.open(file_path) as reader:
                for page in reader.pages:
                    content += page.extract_text() or ""
        elif file_path.lower().endswith(".docx"):
            doc = Document(file_path)
            for para in doc.paragraphs:
                content += para.text + "\n"
        elif file_path.lower().endswith(".md"):  # 新增支持Markdown格式
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
        else:
            print(f"警告：跳过不支持的文件格式 {file_path}")
            return None
    except Exception as e:
        print(f"读取文件 {file_path} 时出错：{e}")
        return None
    return content.strip()


def generate_prompt(document_content):
    """
    glm-4.5模型提示词（基于JSON格式化要求）
    """
    prompt = f"""
    作为一名文档分析助手，善于从文献中提取结构化信息构建知识图谱，请严格基于提供的文献内容，完成以下信息提取任务。
    你的任务是分析文献中描述的甜味剂，并将其相关信息结构化为JSON格式。
    请只返回JSON格式的提取结果，不要包含任何解释性文字、总结或前导语。

    提取规则：
    1.  仔细阅读整篇文献，判断文献研究的致甜成分属于哪一类甜味剂类型，甜味剂类型列表:[天然植物提取物、合成甜味剂、糖类物质、美拉德反应产物、热解反应产物、发酵产物]等
    2.  如果文献中明确提到了某个实体名称，请提取该名称。
    3.  如果没有明显的实体名称，但有相关描述，请用一段文本来代替该实体名称。
    4.  如果文献中没有提及任何相关信息，请使用空字符串 "" 作为该字段的值,不要用未提及、文中未提及等文字填充字段，采用空字符串代""填充该字段。
    5.  输出必须是一个JSON数组。**数组中的每个对象都必须代表一种独立的致甜成分信息**。
    6.  如果文献中包含**多种致甜成分**（例如：烟碱、糠醇），你必须为每种成分创建**一个独立的对象**。
    7.  如果文献中**同一种致甜成分**因为不同的实验条件（例如：不同的添加量、不同的处理方法）而存在**多条独立信息**，你也必须为每条独立信息创建**一个独立的对象**。例如，如果一份文献描述了致甜成分X在添加量为50mg/kg和100mg/kg时的不同余味，你应该创建两个对象，分别记录这两种情况。


    请严格遵循以下JSON结构，并确保所有键名都与我提供的列表完全一致。
    你的输出必须是一个JSON数组，格式如下：[...]。请勿在JSON结构前后添加任何额外文本。

    JSON结构示例：
    [
      {{
        "甜味剂类型": "...",
        "论文名称": "...",
        "致甜成分名称": "...",
        "致甜成分的优点": "...",
        "致甜成分的缺点": "...",
        "致甜成分实验方法": "...",
        "评吸方法"： "...",
        "卷烟余味": "...",
        "致甜成分的应用": "...",
        "致甜成分的添加方法": "...",
        "致甜成分添加量及口感": "...",
        "甜味持久性": "...",
        "致甜成分相对甜度": "...",
        "致甜成分烟气烟丝中的含量及迁移率": "...",
        "致甜成分分析方法": "...",
        "致甜成分的香气": "...",
        "致甜成分的添加部位": "...",
        "烟叶类型": "..."
      }}
    ]

    现在，请参考以下示例，根据提供的文献内容进行提取：

    示例：
    [
      {{
        "甜味剂类型": "天然植物提取物",
        "论文名称": "山奈酚糖苷对卷烟感官品质的改善研究",
        "致甜成分名称": "山奈酚-3-O-β-D-葡萄糖基(1→2)-β-D-半乳糖苷",
        "致甜成分的优点": "能有效改善烟气中的干燥感和刺激感，突出卷烟的甜润感",
        "致甜成分的缺点": "在高温下不稳定，易降解，影响甜味效果",
        "致甜成分实验方法": "",
        "评吸方法": "由10位评价人员以直接嗅闻的方式对各流分的香气特征进行评价，评价结果须7位以上评委意见一致才被接受",
        "卷烟余味": "余味干净、舒适、带有微甜",
        "致甜成分的应用": "可用于卷烟接装纸、滤棒、卷烟纸或直接喷洒于烟丝",
        "致甜成分的添加方法": "将提取物溶解于乙醇后喷洒至烟丝",
        "致甜成分添加量及口感": "添加量为50 mg/kg，可显著提升卷烟的甜感和舒适度",
        "甜味持久性": "甜感温和，持久性较好，无后苦味",
        "致甜成分相对甜度": "蔗糖甜度的2000倍",
        "致甜成分烟气烟丝中的含量及迁移率": "",
        "致甜成分分析方法": "GC-MS分析",
        "致甜成分的香气": "烤甜香、焦甜香",
        "致甜成分的添加部位": "卷烟纸",
        "烟叶类型": "烤烟"
      }}
    ]

    文献内容：
    {document_content}
    """
    return prompt


def call_zhipu_api(prompt):
    """调用智谱AI API获取JSON格式的提取结果"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    data = {
        "model": "glm-4.5",
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.1,
        # "max_tokens": 102400,
        "response_format": {"type": "json_object"}
    }
    try:
        response = requests.post(API_URL, headers=headers, data=json.dumps(data))
        response.raise_for_status()  # 检查HTTP错误
        response_json = response.json()
        if "error" in response_json:
            raise Exception(f"API调用错误：{response_json['error']['message']}")

        content = response_json["choices"][0]["message"]["content"]
        return json.loads(content)
    except requests.exceptions.RequestException as e:
        print(f"请求失败：{e}")
        return None
    except (json.JSONDecodeError, KeyError) as e:
        print(f"解析大模型返回的JSON失败：{e}")
        print("大模型返回的原始内容：", content)
        return None
    except Exception as e:
        print(f"发生未知错误：{e}")
        return None


def write_to_csv(data, file_path, is_first_write):
    """将提取的数据写入CSV文件，is_first_write=True时写入表头"""
    mode = 'w' if is_first_write else 'a'
    with open(file_path, mode, newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=COLUMN_HEADERS)
        if is_first_write:
            writer.writeheader()
        if data:
            writer.writerows(data)


output_dir = os.path.dirname(OUTPUT_CSV_PATH)
if not os.path.exists(output_dir):
    os.makedirs(output_dir, exist_ok=True)


# -------------------------- 主流程 --------------------------
def main():
    if not os.path.exists(DOCUMENTS_DIR):
        print(f"错误：未找到文档文件夹 {DOCUMENTS_DIR}")
        return

    print(f"正在处理文件夹中的所有文档：{DOCUMENTS_DIR}")

    is_first_write = True

    for root, dirs, files in os.walk(DOCUMENTS_DIR):
        for file in files:
            # 新增支持.md文件格式
            if file.lower().endswith(('.pdf', '.docx', '.md')):
                full_path = os.path.join(root, file)

                # 核心改动：在处理单个文件时添加 try-except 块
                try:
                    print(f"--- 正在处理文件：{full_path}")
                    document_content = read_document(full_path)
                    if document_content is None:
                        continue

                    prompt = generate_prompt(document_content)
                    extracted_data = call_zhipu_api(prompt)

                    if extracted_data and isinstance(extracted_data, list):
                        # 补充论文名称
                        paper_name = os.path.splitext(file)[0]
                        valid_data_for_writing = []
                        for entry in extracted_data:
                            if isinstance(entry, dict):
                                entry['论文名称'] = paper_name
                                valid_data_for_writing.append(entry)
                            else:
                                print(f"    -> 警告：文件 '{file}' 中发现无效条目（非字典），已跳过此项。")
                                continue

                        if valid_data_for_writing:
                            write_to_csv(valid_data_for_writing, OUTPUT_CSV_PATH, is_first_write)
                            is_first_write = False
                            print(f"    -> 数据已成功提取并追加至：{OUTPUT_CSV_PATH}")
                        else:
                            print(f"    -> 文件 '{file}' 提取数据为空或所有条目都无效，已跳过。")

                    else:
                        print(f"    -> 提取数据失败或返回格式不正确，跳过文件 {full_path}")

                except Exception as e:
                    # 捕获所有可能的异常，并打印错误信息
                    print(f"    -> 致命错误：处理文件 '{full_path}' 时发生意外错误: {e}")
                    print(f"    -> 已跳过此文件，继续处理下一个文件。")
                    continue  # 继续处理下一个文件

    print("\n所有文档处理完成！")
    print(f"最终结果已合并至：{OUTPUT_CSV_PATH}")


if __name__ == "__main__":
    main()